from typing import List, Optional, Union

from coin.support.data_replay.reader.interface import (
    RawData,
    IRawDataReader)
import coin.flow.archive as archive


class ArchiveReader(IRawDataReader):
  def __init__(self,
               start_timestamp: int,
               end_timestamp: int,
               archive_file: Union[str, List[str]]):
    self._start_timestamp = start_timestamp
    self._end_timestamp = end_timestamp
    self._reader = archive.ArchiveReader.from_file(archive_file)
    self._record_cache = None

  def __del__(self):
    self._reader.close()

  def __enter__(self):
    return self

  def __exit__(self, type, value, traceback):
    self._reader.close()

  def peek(self) -> Optional[RawData]:
    ret = self._record_cache
    if ret is None:
      while True:
        record = self._reader.read()
        if record is None:
          break
        timestamp = record.timestamp
        if timestamp < self._start_timestamp:
          continue
        if timestamp > self._end_timestamp:
          break
        ret = RawData(timestamp=timestamp, value=record.data)
        self._record_cache = ret
        break
    return ret

  def read(self) -> Optional[RawData]:
    record = self._record_cache
    if record is None:
      record = self.peek()
    self._record_cache = None
    return record


if __name__ == '__main__':
  import datetime
  import json
  import google.protobuf.json_format as json_format
  from coin.base.datetime_util import to_datetime, to_timestamp_int
  from coin.exchange.base.strategy_logger import (
    enumerate_logs as enumerate_strat_logs)
  from coin.proto.coin_request_pb2 import StrategyRequestProto
  from coin.proto.coin_strategy_pb2 import StrategyLog

  strat_req = StrategyRequestProto(strategy_name='vmm_sfty_delta_hedge')
  machine = 'strategy-327.ap-northeast-1'
  dt = datetime.date(2023, 6, 13)
  log_dir = '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log/%s/%s' % (
      machine, dt.strftime('%Y%m%d'))
  log_files = enumerate_strat_logs(strat_req, dt, log_dir)
  reader = ArchiveReader(
      to_timestamp_int(datetime.datetime(2023, 6, 13)),
      to_timestamp_int(datetime.datetime(2023, 6, 13, 1)),
      log_files
  )
  datas = []
  prev_ts = 0
  with reader as reader:
    while True:
      record = reader.read()
      if record is None:
        break
      pb = StrategyLog()
      pb.ParseFromString(record.value)
      data = json_format.MessageToDict(pb, preserving_proto_field_name=True)
      curr_ts = record.timestamp
      assert curr_ts >= prev_ts, (prev_ts, curr_ts)
      prev_ts = curr_ts
      hr_ts = to_datetime(curr_ts).strftime('%Y%m%d %H:%M:%S.%f')
      data['record_timestamp'] = curr_ts
      data['record_timestamp_human_readable'] = hr_ts
      datas.append(data)
  records = '[\n' + ',\n'.join(json.dumps(record) for record in datas) + '\n]'
  print(records)
